import tensorflow as tf
import numpy as np

sess = tf.Session()
def a():
    print('simple while_loop')
    i = tf.constant(0)
    c = lambda i:tf.less(i, 10)
    b = lambda i:tf.add(i, 1)
    r = tf.while_loop(c, b, [i])
    print('r:', r)
    print(sess.run((r, i)))

a()

def b():
    print('nest and numetuple while loop')
    import collections
    Pair = collections.namedtuple('Pair', 'j,k')
    ijk_0 = (tf.constant(0), Pair(tf.constant(1), tf.constant(2)))
    c = lambda i, p:i<10
    b = lambda i, p:(i+1, Pair((p.j + p.k), (p.j - p.k)))
    ijk_final = tf.while_loop(c, b, ijk_0)
    print('ijk_final:', sess.run(ijk_final))

b()

def c():
    print('using shape invariants')
    i0 = tf.constant(0)
    m0 = tf.ones([2,2])
    c = lambda i, m: i<10
    b = lambda i, m: [i+1, tf.concat([m, m], axis=0)]
    r = tf.while_loop(c, b, loop_vars=[i0, m0], shape_invariants=[i0.get_shape(), tf.TensorShape([None, 2])])
    print(sess.run(r))

c()

def d():
    import time
    print('non-stric semantics')
    n = 10000
    x = tf.constant(list(range(n)))
    c = lambda i, x: i< n
    b = lambda i, x: (tf.Print(i+1, [i],'i'), tf.Print(x+1, [i], 'x:'))
    i, out = tf.while_loop(c, b, (0, x), parallel_iterations=8)
    start = time.time()
    print(sess.run((out)))
    end = time.time()
    print('', (end - start))

def e():
    print('weird while loop')
    t_var = tf.Variable(0.0)
    print('t_var:', t_var)
    def cond(i, _):
        return tf.less(i, 6)

    def body(i, _):
        #global t_var
        #p_op = tf.Print(t_var, [t_var, i], 't_var:')
        print('t_var in loop:', t_var)
        #with tf.control_dependencies([p_op]):
        flow = tf.assign_add(t_var, tf.constant(1, dtype=tf.float32))
        #flow = tf.assign(t_var, tf.cast(i, tf.float32))
        return tf.add(i, 1), flow

    i, re = tf.while_loop(cond, body, [1, t_var])
    with tf.Session() as sess:
        tf.global_variables_initializer().run()
        print(sess.run([i, re]))
e()